"""
Message Decoder模块 - 修复版
功能：从冗余消息中提取原始消息
这是FakeTagger系统的第5个模块
"""

import torch
import torch.nn as nn
import numpy as np

class MessageDecoder:
    """
    消息解码器 - 修复版
    负责从冗余消息中解码出原始消息，并计算误码率等指标
    """
    
    def __init__(self, message_length=30, redundancy_factor=2):
        """
        初始化消息解码器
        
        Args:
            message_length: 原始消息长度
            redundancy_factor: 冗余因子
        """
        self.message_length = message_length
        self.redundancy_factor = redundancy_factor
        self.redundant_length = message_length * redundancy_factor
        
        print(f"Message Decoder initialized (FIXED):")
        print(f"  - Original message length: {self.message_length}")
        print(f"  - Redundancy factor: {self.redundancy_factor}")
        print(f"  - Expected redundant length: {self.redundant_length}")
        
    def decode_redundant_message(self, redundant_messages):
        """
        从冗余消息中解码原始消息
        修复：使用更鲁棒的解码策略，包括错误纠正
        
        Args:
            redundant_messages: 冗余消息 (batch_size, redundant_length)
            
        Returns:
            torch.Tensor: 解码的原始消息 (batch_size, message_length)
        """
        batch_size = redundant_messages.shape[0]
        device = redundant_messages.device
        
        # 确保输入尺寸正确
        if redundant_messages.shape[1] != self.redundant_length:
            print(f"Warning: Expected redundant length {self.redundant_length}, "
                  f"got {redundant_messages.shape[1]}")
            # 如果长度不匹配，尝试调整
            if redundant_messages.shape[1] > self.redundant_length:
                redundant_messages = redundant_messages[:, :self.redundant_length]
            else:
                # 如果消息太短，进行填充
                padding = torch.zeros(batch_size, 
                                    self.redundant_length - redundant_messages.shape[1],
                                    device=device)
                redundant_messages = torch.cat([redundant_messages, padding], dim=1)
        
        # 将冗余消息重新整形为 (batch_size, redundancy_factor, message_length)
        try:
            reshaped = redundant_messages.view(batch_size, self.redundancy_factor, self.message_length)
        except RuntimeError as e:
            print(f"Reshape error: {e}")
            print(f"Trying to reshape {redundant_messages.shape} to "
                  f"({batch_size}, {self.redundancy_factor}, {self.message_length})")
            # 降级处理：直接截取前面的部分
            decoded = redundant_messages[:, :self.message_length]
            return self._apply_threshold(decoded)
        
        # 修复：使用更鲁棒的解码策略
        if self.redundancy_factor == 1:
            # 无冗余：直接解码
            decoded = reshaped.squeeze(1)
        else:
            # 有冗余：使用多种策略结合
            
            # 策略1：加权平均（给第一份原始消息更高权重）
            weights = torch.ones(self.redundancy_factor, device=device)
            weights[0] = 2.0  # 第一份（原始消息）权重更高
            weights = weights / weights.sum()
            
            weighted_avg = torch.sum(reshaped * weights.view(1, -1, 1), dim=1)
            
            # 策略2：多数投票（对于每个比特位）
            # 先进行二值化
            binary_reshaped = (reshaped > 0.5).float()
            majority_vote = torch.sum(binary_reshaped, dim=1) > (self.redundancy_factor / 2)
            majority_vote = majority_vote.float()
            
            # 策略3：结合策略1和2
            # 如果加权平均和多数投票结果一致，使用加权平均
            # 否则使用多数投票（更可靠）
            weighted_binary = (weighted_avg > 0.5).float()
            consistent_mask = (weighted_binary == majority_vote).float()
            
            # 对于一致的位置，使用加权平均的连续值
            # 对于不一致的位置，使用多数投票的二值结果
            decoded = weighted_avg * consistent_mask + majority_vote * (1 - consistent_mask)
        
        return self._apply_threshold(decoded)
    
    def _apply_threshold(self, messages):
        """
        应用阈值函数，将连续值转换为接近0/1的值
        
        Args:
            messages: 连续值消息
            
        Returns:
            阈值化后的消息
        """
        # 使用软阈值：不完全二值化，但推向0/1
        # 这样既保持了梯度，又增强了二值化特性
        sigmoid_messages = torch.sigmoid((messages - 0.5) * 10)  # 10是锐化因子
        
        # 进一步推向0/1
        thresholded = sigmoid_messages * 0.8 + 0.1  # 映射到[0.1, 0.9]
        
        return thresholded
    
    def calculate_bit_error_rate(self, original_messages, decoded_messages):
        """
        计算误码率 (BER - Bit Error Rate)
        修复：使用更合适的阈值进行比较
        
        Args:
            original_messages: 原始消息 (batch_size, message_length)
            decoded_messages: 解码消息 (batch_size, message_length)
            
        Returns:
            float: 平均误码率 (0-1之间)
        """
        # 二值化比较，使用0.5作为阈值
        original_binary = (original_messages > 0.5).float()
        decoded_binary = (decoded_messages > 0.5).float()
        
        # 计算每个样本的错误比特数
        errors = torch.sum(original_binary != decoded_binary, dim=1)
        total_bits = original_messages.shape[1]
        
        # 计算每个样本的误码率
        sample_bers = errors.float() / total_bits
        
        # 返回平均误码率
        avg_ber = torch.mean(sample_bers)
        
        return avg_ber.item()
    
    def calculate_message_accuracy(self, original_messages, decoded_messages):
        """
        计算消息准确率（完全匹配的消息比例）
        
        Args:
            original_messages: 原始消息
            decoded_messages: 解码消息
            
        Returns:
            float: 消息准确率 (0-1之间)
        """
        original_binary = (original_messages > 0.5).float()
        decoded_binary = (decoded_messages > 0.5).float()
        
        # 检查每个消息是否完全匹配
        perfect_matches = torch.all(original_binary == decoded_binary, dim=1)
        accuracy = torch.mean(perfect_matches.float())
        
        return accuracy.item()
    
    def calculate_soft_bit_error_rate(self, original_messages, decoded_messages):
        """
        计算软误码率（考虑连续值的差异）
        
        Args:
            original_messages: 原始消息
            decoded_messages: 解码消息
            
        Returns:
            float: 软误码率
        """
        # 计算连续值的L1距离
        l1_distance = torch.mean(torch.abs(original_messages - decoded_messages))
        
        # 归一化到[0,1]范围
        soft_ber = l1_distance / 1.0  # 最大可能距离是1
        
        return soft_ber.item()
    
    def decoded_message_error_rate_batch(self, original_messages, decoded_messages):
        """
        批量计算误码率（与原项目兼容的方法）
        
        Args:
            original_messages: 原始消息
            decoded_messages: 解码消息
            
        Returns:
            float: 平均误码率
        """
        return self.calculate_bit_error_rate(original_messages, decoded_messages)

def test_message_decoder():
    """测试消息解码器功能"""
    print("=== Testing Message Decoder (FIXED) ===")
    
    # 导入消息生成器进行联合测试
    from message_generator import MessageGenerator
    
    # 创建生成器和解码器
    generator = MessageGenerator(message_length=30, redundancy_factor=2)
    decoder = MessageDecoder(message_length=30, redundancy_factor=2)
    
    # 测试数据
    batch_size = 4
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    
    # 生成消息
    original_msgs, redundant_msgs = generator.generate_redundant_message(batch_size, device)
    print(f"Generated - Original: {original_msgs.shape}, Redundant: {redundant_msgs.shape}")
    
    # 无噪声解码测试
    decoded_msgs = decoder.decode_redundant_message(redundant_msgs)
    print(f"Decoded messages shape: {decoded_msgs.shape}")
    
    # 计算指标
    ber = decoder.calculate_bit_error_rate(original_msgs, decoded_msgs)
    accuracy = decoder.calculate_message_accuracy(original_msgs, decoded_msgs)
    soft_ber = decoder.calculate_soft_bit_error_rate(original_msgs, decoded_msgs)
    
    print(f"No noise - BER: {ber:.6f}, Message Accuracy: {accuracy:.4f}, Soft BER: {soft_ber:.6f}")
    
    # 添加噪声测试鲁棒性
    print("\n--- Testing with noise ---")
    noise_levels = [0.05, 0.1, 0.2, 0.3, 0.5]
    
    for noise_level in noise_levels:
        noisy_redundant = redundant_msgs + torch.randn_like(redundant_msgs) * noise_level
        noisy_redundant = torch.clamp(noisy_redundant, 0.0, 1.0)  # 限制范围
        decoded_noisy = decoder.decode_redundant_message(noisy_redundant)
        
        ber_noisy = decoder.calculate_bit_error_rate(original_msgs, decoded_noisy)
        acc_noisy = decoder.calculate_message_accuracy(original_msgs, decoded_noisy)
        soft_ber_noisy = decoder.calculate_soft_bit_error_rate(original_msgs, decoded_noisy)
        
        print(f"Noise {noise_level:.2f} - BER: {ber_noisy:.6f}, Accuracy: {acc_noisy:.4f}, Soft BER: {soft_ber_noisy:.6f}")
    
    print("Message Decoder (FIXED) test completed!\n")

if __name__ == "__main__":
    test_message_decoder() 